Summary

One of the key features of the CodemodCommand class is that it abstracts away the final run of AddImportsVisitor and RemoveImportsVisitor in transform_module (Code Source). However, currently developers still need to utilize the static functions:

1. RemoveImportsVisitor.remove_unused_import
2. RemoveImportsVisitor.remove_unused_import_by_node
3. AddImportsVisitor.add_needed_import

To actually manage the import states in the self.context: CodemodContext. I propose defining these three new instance functions on CodemodCommand:

1. remove_unused_import
2. remove_unused_import_by_node
3. add_needed_import

These new lightweight helper wrapper functions defined on CodemodCommand will help to hopefully completely abstract away RemoveImportsVisitor and AddImportsVisitor.

[Example] Remove + Add Imports

Before:

from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor

class RemoveAddImportExampleCommand(VisitorBasedCodemodCommand):
    def visit_Module(self, node: cst.Module) -> None:
        RemoveImportsVisitor.remove_unused_import(self.context, "a.b.c", "bar")
        AddImportsVisitor.add_needed_import(self.context, "hello_module", "bar")

After:

# No more extra imports :)

class RemoveAddImportExampleCommand(VisitorBasedCodemodCommand):
    def visit_Module(self, node: cst.Module) -> None:
        self.remove_unused_import("a.b.c", "bar")  # no need to pass in self.context anymore :)
        self.add_needed_import("hello_module", "bar")

Test Plan

1. Format your code

   uv run poe format

2. Run the type checker

   uv run poe typecheck

3. Test your changes

   uv run poe test

4. Check linters

   uv run poe lint

Specific Unit Tests found in:

libcst/codemod/tests/test_command_helpers.py

and can be run using

uv run poe test -k test_command_helpers

[IG][Meta] Would be nice to have this to lower barrier of entry for using CodemodCommand.